{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#default_exp geometry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Geometry\n",
    "\n",
    "> Geometry transforms and projections. Implemented as Tensorflow layers and differentiable if it makes sense."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "\n",
    "#For an easier life when developing\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import models, layers\n",
    "import tensorflow_addons as tfa\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from radicalsdk.radar import v1_constants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "class PolarToCartesianWarp(layers.Layer):\n",
    "    \"\"\"Differentiable Polar Image to Cartersian Mapping\n",
    "    This is a Tensorflow Keras Layer and\n",
    "    expects a batch of input with shape [n, r, az, c]\n",
    "\n",
    "    Running eagerly is supported as well.\n",
    "    For single example input, use expanddims or newaxis.\n",
    "    \"\"\"\n",
    "    def __init__(self, full=True, scale=1.):\n",
    "        super(PolarToCartesianWarp, self).__init__()\n",
    "        self.full = full\n",
    "        self.scale = scale\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        range_bins, angle_bins = input_shape[1:3]\n",
    "        xx = np.arange(-range_bins, range_bins)/range_bins\n",
    "        yy = 1 - np.arange(0, range_bins)/range_bins\n",
    "\n",
    "        mg = np.meshgrid(xx, yy)\n",
    "\n",
    "        rr = range_bins - np.sqrt(mg[0]**2 + mg[1]**2) * range_bins\n",
    "        tt = angle_bins - np.arctan2(mg[1], mg[0])/np.pi*angle_bins\n",
    "\n",
    "        self.warp = tf.Variable(\n",
    "            np.stack([tt, rr], axis=-1)[np.newaxis, ...],\n",
    "            trainable=False,\n",
    "            dtype=tf.float32\n",
    "        )\n",
    "\n",
    "        self.warp = tf.repeat(self.warp, repeats=input_shape[0], axis=0)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return tfa.image.resampler(inputs, self.warp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "\n",
    "class CameraRadarCoordinateTransform:\n",
    "    \"\"\"Stub Not implemented yet\"\"\"\n",
    "    def __init__(self, camera_f, camera_p, radar_f, radar_p):\n",
    "        self.camera_f = camera_f\n",
    "        self.camera_p = camera_p\n",
    "        self.radar_f = radar_f\n",
    "        self.radar_p = radar_p\n",
    "\n",
    "\n",
    "    def camera2world(self, uv, depth=None):\n",
    "        \"\"\"Projects camera pixel coordinates u, v, depth to world coordinates\n",
    "        uv can be n x 2 lists or n x 3\n",
    "        if uv is n x 2, depth image needs to be provided.\n",
    "        Depth will be retreived using nearest neighbor interpolation\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "    def world2camera(self, xyz):\n",
    "        pass\n",
    "\n",
    "    def radar2world(self, xy, h=None):\n",
    "        pass\n",
    "\n",
    "    def world2radar(self, xyz):\n",
    "        pass\n",
    "\n",
    "    def camera2radar(self, uv, depth=None):\n",
    "        \"\"\"Convenience function combining `camera2world` and `world2radar`\"\"\"\n",
    "        xyz = self.camera2world(uv, depth)\n",
    "        radar_xy = self.world2camera(xyz)\n",
    "\n",
    "        return radar_xy\n",
    "\n",
    "    def radar2camera(self, radar_xy, radar_height=None):\n",
    "        \"\"\"Convenience function combining `radar2world` and `world2camera`\"\"\"\n",
    "        xyz = self.radar2world(radar_xy, radar_height)\n",
    "        uvz = self.world2camera(xyz)\n",
    "\n",
    "        return uvz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def compute_radar_intrinsic_matrix(radarframe):\n",
    "    \"\"\"Radar frame needs to provide max_range and range_nbins\"\"\"\n",
    "    scale = 1/radarframe.max_range\n",
    "    nbins = radarframe.range_nbins\n",
    "    if radarframe.flipped:\n",
    "        f = np.array([\n",
    "            [nbins, 0,      nbins],\n",
    "            [0,     -nbins, nbins],\n",
    "            [0,     0,      1./scale]\n",
    "        ])\n",
    "    else:\n",
    "        f = np.array([\n",
    "            [nbins, 0,      nbins],\n",
    "            [0,     -nbins, nbins],\n",
    "            [0,     0,      1/scale]\n",
    "        ])\n",
    "\n",
    "    return scale * f"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Get radar intrinsics from radar config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[ 21.3503495   0.         21.3503495]\n",
      " [  0.        -21.3503495  21.3503495]\n",
      " [  0.          0.          1.       ]]\n"
     ]
    }
   ],
   "source": [
    "from radicalsdk.radar.config_v1 import read_radar_params\n",
    "from radicalsdk.radar.v1 import RadarFrame\n",
    "\n",
    "radar_config = read_radar_params('../samples/indoor_human_rcs.cfg')\n",
    "rf = RadarFrame(radar_config)\n",
    "\n",
    "radar_f = compute_radar_intrinsic_matrix(rf)\n",
    "print(radar_f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
